package com.heima.leetcode.practice;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

/**
 * 字典树
 * @author 勾新杰
 * @version 1.0
 * @date 2024/10/27 11:24
 */
public class Trie {
    private final TrieNode root; // 根节点，存放的空字符

    /**
     * 初始化字典树
     */
    public Trie() {
        root = new TrieNode();
    }

    /**
     * 字典树节点
     */
    static class TrieNode {
        TrieNode[] children; // 默认放的是26个小写字母
        boolean isEnd; // 是否是某个单词的结尾

        /**
         * 构造方法，初始化26个孩子，以及是否是结尾
         */
        public TrieNode() {
            this.children = new TrieNode[26];
            this.isEnd = false;
        }

        /**
         * 判断当前节点是否只有一个孩子
         * @return true：只有一个孩子，false：否则
         */
        public boolean hasOnlyChild(){
            TrieNode onlyChild = null;
            for (TrieNode child : children) {
                if (child != null){
                   if (onlyChild == null){
                       onlyChild = child;
                   }else {
                       return false;
                   }
                }
            }
            return onlyChild != null;
        }
    }

    /**
     * 插入操作
     * @param word 单词，字符串
     */
    public void insert(String word){
        TrieNode curr = root;
        char[] array = word.toCharArray();
        for (char c : array) {
            int i = c - 'a';
            if (curr.children[i] == null) {
                curr.children[i] = new TrieNode();
            }
            curr = curr.children[i];
        }
        curr.isEnd = true;
    }

    /**
     * 查找操作
     * @param word 带查找单词
     * @return 查找到到返回true，否则返回false
     */
    public boolean search(String word){
        TrieNode curr = root;
        char[] array = word.toCharArray();
        for (char c : array) {
            int index = c - 'a';
            if (curr.children[index] != null){
                curr = curr.children[index];
            }else {
                return false;
            }
        }
        return curr.isEnd;
    }

    /**
     * 前缀匹配
     * @param prefix 前缀
     * @return 匹配到返回true，否则返回false
     */
    public boolean startsWith(String prefix){
        TrieNode curr = root;
        char[] array = prefix.toCharArray();
        for (char c : array) {
            int index = c - 'a';
            if (curr.children[index] != null){
                curr = curr.children[index];
            }else {
                return false;
            }
        }
        return true;
    }

    /**
     * 删除操作
     * @param word 待删除单词
     */
    public void delete(String word){
        TrieNode curr = root;
        char[] array = word.toCharArray();
        LinkedList<TrieNode> nodeStack = new LinkedList<>();
        LinkedList<Integer> indexStack = new LinkedList<>();
        for (char c : array) {
            int index = c - 'a';
            if (curr.children[index] != null){
                nodeStack.push(curr);
                indexStack.push(index);
                curr = curr.children[index];
            }else {
                throw new IllegalArgumentException(word + "不存在");
            }
        }
        if (!curr.isEnd){
            throw new IllegalArgumentException(word + "不存在");
        }
        while (!nodeStack.isEmpty()){
            TrieNode pop = nodeStack.pop();
            int index = indexStack.pop();
            if (!pop.hasOnlyChild()){
                pop.children[index] = null;
                return;
            }
        }
    }

    /**
     * 自动补全
     * @param prefix 前缀
     * @return 补全后的所有单词
     */
    public List<String> autoComplete(String prefix){
        TrieNode curr = root;
        char[] array = prefix.toCharArray();
        for (char c : array) {
            int index = c - 'a';
            if (curr.children[index] != null){
                curr = curr.children[index];
            }else {
                return new ArrayList<>();
            }
        }
        return getAllWords(curr, prefix);
    }

    /**
     * 递归获取所有单词
     * @param curr 当前节点
     * @param prefix 前缀
     * @return 所有单词
     */
    private List<String> getAllWords(TrieNode curr, String prefix) {
        List<String> list = new ArrayList<>();
        if (curr.isEnd){
            list.add(prefix);
        }
        for (int i = 0; i < curr.children.length; i++) {
            TrieNode child = curr.children[i];
            if (child != null){
                list.addAll(getAllWords(child, prefix + (char)('a' + i)));
            }
        }
        return list;
    }

    /**
     * 当前字典树的最长前缀，就是当前字典树存放的这些字符串的最长公共前缀
     * @return 最长公共前缀
     */
    public String longestPrefix(){
        TrieNode curr = root;
        StringBuilder builder = new StringBuilder();
        // 当curr没到一个单词的结尾时，继续往下，如果中间有分叉，直接找到公共前缀，返回即可
        while (!curr.isEnd){
            int onlyChildIndex = -1; // 唯一孩子的索引
            for (int i = 0; i < curr.children.length; i++) {
                TrieNode child = curr.children[i]; // 当前孩子
                if (child != null) { // 不是null，说明有孩子
                    if (onlyChildIndex == -1) { // 唯一孩子还没有，则加上
                        onlyChildIndex = i;
                    } else { // 唯一孩子有了，则说明不止一个孩子，有分支，返回结果
                        return builder.toString();
                    }
                }
            }
            if (onlyChildIndex != -1){ // 找到了唯一孩子，则加入到字符串中
                builder.append((char) ('a' + onlyChildIndex)); // 这里记得强转
                curr = curr.children[onlyChildIndex]; // 别忘了更新curr
            }
            // 没找唯一孩子，说明当前是单词的结束位，while自己会退出循环
        }
        return builder.toString();
    }
}
